import os
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
import zipfile
from utils.preprocessing import *

rawdata_root = '/media/gserver/data/FashionAI'


mode = 'online'


merge_list = [
    'res101-[0567]-9179-aug2-9243(dilate-pred).csv',
    'res101-[1234]-9623-aug2-9677(dilate-pred).csv'
]
# model_name = 'res101-[1234]-0.9677-[0567]-0.9243-aug2(dilate-pred).csv'
model_name = 'res101-[1234]-0.9677-[0567]-0.9243-aug2(dilate-pred)'

if mode == 'online':
    csv_root = './online_pred2/part_csv/'
    test_root = os.path.join(rawdata_root, 'rank')
    test_pd = pd.read_csv(os.path.join(rawdata_root, 'round2/week-rank/Tests/question.csv'),
                       header=None, names=['ImageName', 'AttrKey', 'AttrValues'])

else:

    csv_root = './val_pred2/part_csv/'
    round1_df = pd.read_csv(os.path.join(rawdata_root, 'round1/base/Annotations/label.csv'),
                            header=None, names=['ImageName', 'AttrKey', 'AttrValues'])
    round1_df = join_path_to_df(round1_df, rawdata_root, 'round1/base')

    round2_df = pd.read_csv(os.path.join(rawdata_root, 'round2/train/Annotations/label.csv'),
                            header=None, names=['ImageName', 'AttrKey', 'AttrValues'])
    round2_df = join_path_to_df(round2_df, rawdata_root, 'round2/train')

    extra_df = pd.read_csv(os.path.join(rawdata_root, 'testab_warmup_for_train.txt'),
                           header=None, names=['ImageName', 'AttrKey', 'AttrValues', 'hash']).drop(['hash'], axis=1)
    extra_df = join_path_to_df(extra_df, rawdata_root, 'round1/testa_b')

    round2_train_pd, test_pd = train_test_split(round2_df, test_size=0.1, random_state=37,
                                               stratify=round2_df['AttrKey'])

    train_pd = pd.concat([round2_train_pd, round1_df, extra_df], axis=0, ignore_index=True)
    train_pd.index = range(train_pd.shape[0])

test_pd = test_pd.drop('AttrValues', axis=1)

concat_pd = pd.DataFrame()
for file_name in merge_list:
    file_path = os.path.join(csv_root, file_name)
    part_pd = pd.read_csv(file_path,header=None, names=['ImageName', 'AttrKey', 'AttrValueProbs'])

    concat_pd = pd.concat([concat_pd, part_pd],axis=0,ignore_index=True)

concat_pd.index = range(concat_pd.shape[0])

test_pd = test_pd.merge(concat_pd, on=['ImageName', 'AttrKey'], how='left')
print test_pd.head()
print test_pd.info()


# make zip file
if mode == 'online':
    test_pd[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('online_pred2/csv/%s.csv' % model_name, header=None,
                                                               index=False)

    z = zipfile.ZipFile('./online_pred2/subs_zip/%s.zip'%model_name, 'w', zipfile.ZIP_DEFLATED)
    z.write('./online_pred2/csv/%s.csv' % model_name, arcname='%s.csv' % model_name)
    z.close()

else:
    test_pd[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('val_pred2/csv/%s.csv' % model_name, header=None,
                                                               index=False)
    # todo ADD MAP EVAL WHEN CONCAT VAL


